CPVD Cross Project Vulnerability Detection Based on Graph Attention Network and Domain Adaptation

0 Abstract

代码漏洞检测对于软件安全预防至关重要。大规模软件代码中的漏洞注释非常繁琐且具有挑战性,这需要领域专家花费大量时间进行注释。这项工作提供了CPVD,这是一种跨域漏洞检测方法**,基于“学习使用一个具有丰富漏洞标签的项目快速预测另一个项目的漏洞标签”的挑战CPVD使用代码属性图来表示代码,并使用图注意力网络和卷积池网络来提取图特征向量。**在跨域漏洞检测的域自适应表示学习阶段,它减少了源域和目标域数据之间的分布。在本文中,我们在不同的真实世界项目代码上相互测试。与没有域自适应的方法和基于自然语言处理的域自适应方法相比,CPVD更通用,在跨域漏洞检测任务中表现更好。具体而言,对于chr_deb、qemu、libav和sard这四个数据集,它们的F1得分分别为70.2%、81.1%、59.7%和78.1%,AUC分别为88.4%、86.3%、85.2%和88.6%。

代码属性图,跨域漏洞检测,域自适应表示学习,图注意力网络。

1 Intro or Overview

1.1 Problem and Challenge

在代码漏洞检测任务中,VulDeePecker、μVulDeePecker、SySeVR、Vuldeeplocator、Devign、BGNN4VD、Reveal和Ivdetect已经表明,使用神经网络进行自动特征提取比专家制作的特征具有更好的性能。VulDeePecker、μVulDeePecker、SySeVR和Vuldeeplocator将代码函数处理为标记序列,标记序列被处理为自然语言文本。然而,Devgin、BGNN4VD、Reveal和Ivdetect体现了通过图神经网络(以下简称GNN)提取代码函数的图结构特征。这些方法已被证明优于特征提取方法,如递归神经网络、Bi-LSTM和GRU。

然而,前面的技术都导致了漏洞识别问题中的另一个重要问题:项目中缺乏易感代码标签。数据集及其标签用于推动当前的深度学习模型。深度学习模型的预测性能由数据集的数量和质量以及它们的标记决定。由于漏洞标签的稀缺性,历史漏洞不足以训练和验证神经网络模型,尤其是对处于休眠状态的开源项目。

Vulnerability detection in large-scale software code is timeconsuming, complicated, and error-prone;

尽管源域和目标属于不同项目的漏洞代码集,但它们在相同的特征提取器后具有相似的特征空间和标签分布。尽管如此,它们的概率联合分布在跨域漏洞检测问题上更进一步。源域和目标域数据集用于漏洞分类任务,因此最终的分类目标是相同的。基于上述前提条件,可以在跨域漏洞识别中使用域自适应方法。

1.2 Motivation

1.3 Contribution

总之,本文的贡献如下:本文提出了一种将图注意力网络和域自适应表示学习相结合的跨域漏洞检测方法CPVD。这是图神经网络与领域自适应相结合进行跨领域漏洞检测的开端。基于这一想法,研究人员可以提出不同的方法来提高漏洞检测性能。本文验证了域自适应方法更适合于未标记的漏洞检测任务。本文验证了在跨域漏洞检测任务中,代码的图形表示优于令牌序列处理。本文验证了只有对源域进行重新采样才能提高漏洞检测性能。

2 Architecture & Method

2.1 System Overview

image-20240313154641145

问题定义

跨域漏洞代码检测是一个二进制分类问题,旨在将目标域代码分为易受攻击和不易受攻击。跨域漏洞代码检测有两种域分布,即源域代码分布S(C,y)和目标域代码分布T(C,?),其中C是代码函数;y表示漏洞分类标签,y∈{0,1},0表示没有漏洞,1表示有漏洞,“?”表示未知标签。此外,这两个域分布都有域标签,d∈{S,T},如果d=S,则xG~S(xG);else xG~T(xG),其中xG是样本的图特征向量。跨域漏洞检测的目标是训练一个神经网络Nf(C),以不断减少源域漏洞分类损失和域分类损失,最终实现对目标域中的代码漏洞进行准确分类的目标。上述损失可以正式定义为

image-20240313155919021

其中,L是源域漏洞分类损失,y是漏洞分类标签,Ld是域分类损失,d是域标签。

代码预处理

image-20240314212644742

每个节点都包含一个键、代码语句和属性元素(例如Identifier、AssignmentExpression、ParameterType、ExpressionStatement),边表示节点之间的关系,边类型是对它们关系的描述,类型为IS_AST_PARENT、FLOWS_TO、DEF、USE、CONTROLS等。注意,每个函数都有不同的CPG,因为它们的语义和句法结构不同,所以函数的代码属性图不一定包含所有的边类型。改论文中使用了10种类型的边。

节点和边的类型分别用one-hot表示,语句用词嵌入技术进行表示。

图特征提取

图特征提取阶段的输入是标记的源域节点向量和未标记的目标域节点向量;对于标记的源域图,在预训练后输出图的特征向量,而对于未标记的目标域图,使用在源域中训练的模型来提取目标域图的向量特征。

代码预处理阶段输出的节点特征向量是独立于其他节点获得的,因此节点信息较差。代码属性图是根据代码之间的句法和语义结构构建的,每个节点都有一个语句片段,相邻节点之间存在很大的相关性和依赖性。为了最佳地表示节点特征,有必要将其相邻节点的信息映射到自身。因此使用具有双头注意力机制的图注意力网络。

由于源域代码具有标签,我们可以根据漏洞分类任务进行预训练,以获得源域图特征向量和训练后的模型。目标域图特征向量可以从训练的模型中获得。预训练损失函数是一个二分类交叉熵损失函数。

领域自适应表示学习阶段

领域自适应表示学习阶段由四个部分组成:重采样和特征映射、源领域漏洞分类器和领域分类器。

image-20240314215125905

Resampling and Feature Mapping

重采样方法使用SMOTETopek,它是过采样和欠采样的组合。重新采样后的源域平衡数据集和目标域不平衡数据集将进入表示学习网络。

Source Domain Vulnerability Classifier

漏洞分类器Cy(xL,yi)的输入是源域代码和源域漏洞标签的特征向量。我们使用完全连接层作为源域漏洞分类器。在每个完全连接的层之后,Relu被用作激活函数,Dropout被用于防止过拟合。为了使目标域样本接近源域样本,除了使用分类损失函数外,我们还设计了域自适应损失函数LST,以不断减少源域和目标域之间的分布差异。

image-20240314222420388

image-20240314222501665

image-20240314222647564

image-20240314222753871

Domain Classifier

在域分类器中,源域代码的标签为S,目标域代码的标记为T。因此,为了混淆源域和目标域,我们需要最大化域分类误差。域分类器Cd(xG,di)的输入是源或目标域码及其域标签,也就是说,xL∈SõT。我们使用全连接层作为域分类器,在每个全连接层之后使用Relu作为激活函数,并使用Dropout来防止过拟合。DANN[42]设计了一个梯度反转层,确保在反向传播过程中梯度方向自动反转,并在正向传播中进行身份转换。

image-20240314222811899

image-20240314222820817

领域自适应表示学习阶段利用领域数据分布自适应的思想,在训练过程中训练整个领域表示学习网络。因此,模型训练有两个目标:第一是减少代码漏洞分类错误,以确保源域数据的正确分类;第二是增加域分类错误,混淆两个域的代码输入。因此,这一阶段的总损失函数包括两个部分:源域漏洞分类损失和域分类损失。

3 Experiment and Evaluation

与没有域自适应的漏洞检测方法相比,CPVD在漏洞检测任务中的表现如何?

与适用于领域的漏洞检测方法相比,CPVD在漏洞检测任务中的表现如何?

本文中的图特征向量提取阶段的设计如何影响漏洞检测的性能?

对源域数据重新采样如何影响目标域中的漏洞检测性能?

与最先进的领域自适应方法相比,我们采用的领域自适应表示学习方法如何影响漏洞检测任务?

3.1 DataSet and Process

3.2 Evaluation

4 Conclusion

Summary